from typing import Dict, Any

from kafka import KafkaProducer
from kafka.errors import NoBrokersAvailable, kafka_errors

from errors import ErrorCounterException
from flow import Flow
from pyshark_analyzer import PySharkAnalyzer
from utils import kafka_topics
from vpn_analyzers import PPTPAnalyzer, GREAnalyzer, L2TPAnalyzer, PPPAnalyzer, OpenVPNAnalyzer


def send_flows(filename: str):
    ps_analyzer = PySharkAnalyzer(file=filename)
    ps_analyzer.add_vpn_analyzers(
        PPTPAnalyzer(),
        GREAnalyzer(),
        L2TPAnalyzer(),
        PPPAnalyzer(),
        OpenVPNAnalyzer(),
    )
    try:
        producer = KafkaProducer(
            bootstrap_servers=["10.26.22.229:9092"],
            key_serializer=lambda k: k.encode(),
            value_serializer=lambda v: v.encode(),
        )
    except NoBrokersAvailable:
        print("No broker available. plesase check the defined brokers' IP and port")
        exit(1)

    err_cnt = 0

    for flow in ps_analyzer.analyze_packet():  # type: Flow
        print("---------------------------------------------------")
        for vpn_protocol in flow.vpn_protocols:
            vpn_topic = kafka_topics.get(vpn_protocol, None)
            if not vpn_topic:
                continue

            print(f"flow index: {flow.stream_index}\n"
                  f"flow transport type: {flow.flow_transport_type}\n"
                  f"flow recv time: {flow.recv_time}\n"
                  f"flow time: start at {flow.start_time}, end at {flow.end_time}\n"
                  f"flow src IP: {flow.src_ip}, flow dst IP: {flow.dst_ip}\n"
                  f"flow src port: {flow.src_port}, flow dst port: {flow.dst_port}\n"
                  f"flow tunnel's src IP: {flow.tunnel_src_ip}, flow tunnel's dst IP: {flow.tunnel_dst_ip}\n"
                  f"flow tunnel's src port: {flow.tunnel_src_port}, flow tunnel's dst port: {flow.tunnel_dst_port}\n"
                  f"flow vpn protocols: {flow.vpn_protocols}\n"
                  f"flow bidirectional packet length: {flow.bidirectional_packet_length}\n"
                  f"flow bidirectional packet num: {flow.bidirectional_packet_num}\n"
                  f"flow c2s packet length: {flow.c2s_packet_length}, flow c2s packet num: {flow.c2s_packet_num}\n"
                  f"flow s2c packet length: {flow.s2c_packet_length}, flow s2c packet num: {flow.s2c_packet_num}\n"
                  f"flow URL: {flow.url}, flow SNI: {flow.sni}")

            for field in flow.udps.field_names:
                protocol_data = getattr(flow.udps, field)  # type: Dict[str, Any]
                print(f"{field}:")
                for key, value in protocol_data.items():
                    print(f"{key}: {value}")

            data = str(flow)
            print(f"Try to send flow {flow.flow_transport_type} {flow.stream_index} to topic {vpn_topic}.")
            try:
                future = producer.send(topic=vpn_topic,
                                       key=f"{flow.flow_transport_type} {flow.stream_index}",
                                       value=data,
                                       partition=1)

                future.get(timeout=5)

                if not future.succeeded():
                    print(f"flow {flow.stream_index} with type {flow.flow_transport_type} failed to send to kafka")
            except kafka_errors as e:
                print(e)
                err_cnt += 1
                if err_cnt > 1000:
                    raise ErrorCounterException(f"Errors happened in kafka more than {err_cnt} times.")

            break  # stop traversing other vpn protocols in the same flow

    print("All the flows have been analyzed and sent, exited")


if __name__ == "__main__":
    while True:
        send_flows("./data/l2tp-pptp-ipsec/t1-test/l2tp_ctrl_data_full.pcap")
        send_flows("./data/l2tp-pptp-ipsec/t1-test/openvpn_onestream.pcap")
        send_flows("./data/l2tp-pptp-ipsec/t1-test/pptp_clear.pcap")
